import numpy as np
from scipy.io import loadmat
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

# original_BFM = loadmat("3DMM/01_MorphableModel.mat")
original_BFM = loadmat("data_utils/face_tracking/3DMM/01_MorphableModel.mat")
# sub_inds = np.load("3DMM/topology_info.npy", allow_pickle=True).item()["sub_inds"]

shapePC = original_BFM["shapePC"]
shapeEV = original_BFM["shapeEV"]
shapeMU = original_BFM["shapeMU"]
texPC = original_BFM["texPC"]
texEV = original_BFM["texEV"]
texMU = original_BFM["texMU"]

## 可视化平均形状 (shapeMU)
# 假设 shapeMU 是一个 1D 数组，需要重塑为 (n_vertices, 3)
vertices = shapeMU.reshape(-1, 3)

fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
ax.scatter(vertices[:, 0], vertices[:, 1], vertices[:, 2], c='b', marker='.')
# plt.show()

## 可视化形状主成分 (shapePC)
# 选择前几个主成分进行可视化
n_components = 5

for i in range(n_components):
    pc = shapePC[:, i].reshape(-1, 3)
    
    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d')
    ax.scatter(pc[:, 0], pc[:, 1], pc[:, 2], c='r', marker='.')
    plt.title(f"Shape Principal Component {i+1}")
    # plt.show()
    
# 假设 texMU 是一个 1D 数组，需要重塑为 (n_vertices, 3) 用于 RGB 颜色
tex_colors = texMU.reshape(-1, 3)

# 归一化颜色值到 [0, 1] 范围
tex_colors = (tex_colors - tex_colors.min()) / (tex_colors.max() - tex_colors.min())

fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
ax.scatter(vertices[:, 0], vertices[:, 1], vertices[:, 2], c=tex_colors, marker='.')
# plt.show()

## 可视化纹理平均值 (texMU)
# 假设 texMU 是一个 1D 数组，需要重塑为 (n_vertices, 3) 用于 RGB 颜色
tex_colors = texMU.reshape(-1, 3)

# 归一化颜色值到 [0, 1] 范围
tex_colors = (tex_colors - tex_colors.min()) / (tex_colors.max() - tex_colors.min())

fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
ax.scatter(vertices[:, 0], vertices[:, 1], vertices[:, 2], c=tex_colors, marker='.')
# plt.show()

## 可视化特征值 (shapeEV 和 texEV)
plt.figure(figsize=(10, 5))
plt.subplot(121)
plt.plot(shapeEV)
plt.title("Shape Eigenvalues")
plt.xlabel("Principal Component")
plt.ylabel("Eigenvalue")

plt.subplot(122)
plt.plot(texEV)
plt.title("Texture Eigenvalues")
plt.xlabel("Principal Component")
plt.ylabel("Eigenvalue")

plt.tight_layout()
# plt.show()

## 使用 PyVista 进行更高级的 3D 可视化
import pyvista as pv

# 创建 PolyData 对象
mesh = pv.PolyData(vertices)

# 添加颜色数据
mesh.point_data["colors"] = tex_colors

# 创建 Plotter 对象并显示网格
plotter = pv.Plotter()
plotter.add_mesh(mesh, scalars="colors", rgb=True)
# plotter.show()

## 可视化形状变化
def visualize_shape_variation(shapeMU, shapePC, shapeEV, component=0, scale=3, texMU=None, texPC=None, texEV=None):
    mean_shape = shapeMU.reshape(-1, 3)
    pc = shapePC[:, component].reshape(-1, 3)
    ev = shapeEV[component]

    variation = mean_shape + scale * np.sqrt(ev) * pc

    fig = plt.figure(figsize=(12, 5))
    
    if texMU is None:
        ax1 = fig.add_subplot(121, projection='3d')
        ax1.scatter(mean_shape[:, 0], mean_shape[:, 1], mean_shape[:, 2], c='b', marker='.')
        ax1.set_title("Mean Shape")
    else:
        tex_colors = texMU.reshape(-1, 3)
        # 归一化颜色值到 [0, 1] 范围
        tex_colors = (tex_colors - tex_colors.min()) / (tex_colors.max() - tex_colors.min())
        ax1 = fig.add_subplot(121, projection='3d')
        ax1.scatter(mean_shape[:, 0], mean_shape[:, 1], mean_shape[:, 2], c=tex_colors, marker='.')
        ax1.set_title("Mean Shape")

    if texMU is None:
        ax2 = fig.add_subplot(122, projection='3d')
        ax2.scatter(variation[:, 0], variation[:, 1], variation[:, 2], c='r', marker='.')
        ax2.set_title(f"Shape Variation (PC {component+1})")
    else:
        tex_colors = texMU.reshape(-1, 3) + scale * np.sqrt(texEV[component]) * texPC[:, component].reshape(-1, 3)
        # 归一化颜色值到 [0, 1] 范围
        tex_colors = (tex_colors - tex_colors.min()) / (tex_colors.max() - tex_colors.min())
        ax2 = fig.add_subplot(122, projection='3d')
        ax2.scatter(variation[:, 0], variation[:, 1], variation[:, 2], c=tex_colors, marker='.')
        ax2.set_title(f"Shape Variation (PC {component+1})")

    plt.tight_layout()

# visualize_shape_variation(shapeMU, shapePC, shapeEV, component=1)
visualize_shape_variation(shapeMU, shapePC, shapeEV, 1, 1, texMU, texPC, texEV,)
plt.show()